import os
import argparse
import cv2
import numpy as np
from PIL import Image
# import satellite_cloud_generator as scg
import torch

parser = argparse.ArgumentParser()

parser.add_argument('--interval', type=int, default=3)
parser.add_argument('--light', type=float, default=0)
parser.add_argument('--fog', type=float, default=0)
parser.add_argument('--noise', type=float, default=0)
parser.add_argument('--geo', type=float, default=0)
parser.add_argument('--scale', type=float, default=0)
args = parser.parse_args()

input_folder = os.path.join(os.getcwd(), 'input_images')
output_folder = os.path.join(os.getcwd(), 'output_images')
if not os.path.exists(input_folder):
    os.makedirs(input_folder)
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

del_list = os.listdir(output_folder)
for f in del_list:
    file_path = os.path.join(output_folder, f)
    if os.path.isfile(file_path):
        os.remove(file_path)

input_images = os.listdir(input_folder)

for index, image_file in enumerate(input_images):
    if index % args.interval != 0:
        continue
    input_path = os.path.join(input_folder, image_file)
    output_path = os.path.join(output_folder, image_file)
    image = cv2.imread(input_path)
    # 调整亮度
    if args.light >= 1 and args.light <= 99:
        # 调整亮度，对比度系数设置为1表示不改变对比度
        image = cv2.convertScaleAbs(image, alpha=1, beta=args.light)
        # cv2.imwrite(output_path, image)
    if args.geo >= 1 and args.geo <= 100:
        w, h = image.shape[1], image.shape[0]
        k = np.zeros((3,3))
        k[0, 0] = w
        k[1, 1] = h
        k[0, 2] = w / 2.0
        k[1, 2] = h / 2.0
        k[2, 2] = 1.0
        dist_coef = np.zeros((4,1))
        dist_coef[0,0] = args.geo / 100
        image = cv2.undistort(image, k, dist_coef)
    if args.scale >= 1 and args.scale <= 100:
        # 读取原始图像
        scale_percent = 100 - args.scale

        # 获取原始图像的宽度和高度
        original_width = image.shape[1]
        original_height = image.shape[0]

        # 计算缩放后的新宽度和高度
        new_width = int(original_width * scale_percent / 100)
        new_height = int(original_height * scale_percent / 100)

        # 使用OpenCV的resize函数来调整图像大小
        image = cv2.resize(image, (new_width, new_height))
    cv2.imwrite(output_path, image)

    # if args.fog >= 1 and args.fog <= 99:
    #     image = np.array(Image.open(output_path))
    #     image_tensor = torch.from_numpy(np.transpose(image, (2, 0, 1))).unsqueeze(0).float() / 255.0
    #     cloudy_image, cloud_mask = scg.add_cloud(image_tensor, min_lvl=0.0, max_lvl=args.fog / 99, return_cloud=True)
    #     result_image = np.transpose(cloudy_image.squeeze().numpy(), (1, 2, 0))
    #     im = Image.fromarray(np.uint8(result_image * 255))
    #     im.save(output_path)
    if args.noise >= 1 and args.noise <= 99:
        # 打开原始图像
        image = Image.open(output_path)
        # 将图像转换为NumPy数组
        image_np = np.array(image)
        # 生成与图像相同形状的高斯噪声
        mean = 0
        std = args.noise / 2
        noise = np.random.normal(mean, std, image_np.shape).astype(np.uint8)
        # 添加噪声到图像
        noisy_image_np = np.clip((image_np + noise), 0, 255).astype(np.uint8)
        # 创建添加噪声后的图像对象
        noisy_image = Image.fromarray(noisy_image_np)
        # 保存添加噪声后的图像
        noisy_image.save(output_path)